import os
import json
import sys
from tqdm import tqdm

def convert_data(input_json_file, output_path, group_num):
    # 读取 JSON 文件
    with open(input_json_file, "r", encoding="utf-8") as f:
        data = json.load(f)

    # 指令模板
    # instruction_template = "指令：假设你是一个空对空作战经验丰富的空军飞行员，能根据战场态势做出最正确的机动决策。"
    qa_pairs = []

    # 遍历数据并格式化
    # qa_pairs.append(f"{instruction_template}\n")
    for item in tqdm(data, desc="Processing data"):
        question = f"问：{item['instruction']}"
        answer = f"答：{item['output']}"
        qa_pairs.append(f"{question}\n{answer}\n")

    # 确保输出目录存在
    os.makedirs(output_path, exist_ok=True)

    # 写入到新文件
    output_file = os.path.join(output_path, "Air_Combat_Dataset_LLM.txt")
    epoch_num = int(group_num)
    with open(output_file, "w", encoding="utf-8") as f:
        for i in range(len(qa_pairs)):
            # 写入当前 QA
            f.write(qa_pairs[i])
            # 每 epoch_num 条数据后添加空行
            if (i + 1) % epoch_num == 0:
                f.write("\n")  # 添加两行空行作为分隔
    print(f"数据已转换并保存到 {output_file}")

if __name__ == "__main__":
    if len(sys.argv) != 4:
        print("用法: python script.py <输入JSON文件> <输出路径>")
        sys.exit(1)

    input_json_file = sys.argv[1]
    output_path = sys.argv[2]
    group_num = sys.argv[3]
    convert_data(input_json_file, output_path, group_num)
